import torch
import torch.nn as nn
import yaml
import os
from light_rl.utils.dict_wrapper import DictWrapper

import datetime
import numpy as np


def set_seed(seed: int) -> None:
    """
    设置随机种子
    :param seed: int 随机种子
    :return: None
    """
    if seed == -1:
        return
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)


def get_device(device: str) -> torch.device:
    """
    获取设备
    :param device: str 设备名称
    :return: torch.device 设备
    """
    if device.startswith("cpu"):
        return torch.device("cpu")
    # 如果是cuda，判断是否可用
    if device.startswith("cuda"):
        if torch.cuda.is_available():
            return torch.device(device)
        else:
            print("CUDA is not available, use CPU instead.")
            return torch.device("cpu")
    # 如果是mps，判断是否可用
    if device.startswith("mps"):
        if torch.has_mps:
            return torch.device(device)
        else:
            print("MPS is not available, use CPU instead.")
            return torch.device("cpu")


def get_now_time():
    now_time = datetime.datetime.now().strftime("%Y-%m-%d %Hh%Mm%Ss")
    return now_time


def load_config(config_file_path: str) -> DictWrapper:
    """
    加载配置文件
    :param config_file_path: str 配置文件路径
    :return:
    """
    # 加载YAML文件并解析为字典
    with open(config_file_path, 'r') as config_file:
        config = yaml.safe_load(config_file)
    # 将字典转换为DictWrapper对象
    config = DictWrapper(config)
    return config


def get_absolute_path(file: str, relative_path: str) -> str:
    """
    获取相对路径的绝对路径
    :param file: str 当前文件的绝对路径
    :param relative_path: str 相对路径
    :return: str 绝对路径
    """
    # 获取相对路径的绝对路径
    return os.path.join(os.path.dirname(file), relative_path)


def build_mlp(dims: [int]) -> nn.Sequential:  # MLP (MultiLayer Perceptron)
    net_list = []
    for i in range(len(dims) - 1):
        net_list.extend([nn.Linear(dims[i], dims[i + 1]), nn.ReLU()])
    del net_list[-1]  # remove the activation of output layer
    return nn.Sequential(*net_list)


if __name__ == '__main__':
    device_str = "cuda"
    device = torch.device(device_str)
    if device_str == "cuda":
        print(torch.cuda.is_available())
        print(torch.cuda.device_count())
        print(torch.cuda.get_device_name(0))
    print(device)
